package com.shen.langchain4j.config;

import com.shen.langchain4j.service.ChatAssistant;
import com.shen.langchain4j.service.ChatMemoryAssistant;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiTokenCountEstimator;
import dev.langchain4j.service.AiServices;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * 大模型配置类
 */
@Configuration
public class LLMConfig {
    /**
     * 普通对话模型配置
     *
     * @return chatLanguageModel
     */
    @Bean(value = "qwen")
    public ChatModel chatLanguageModelQwen() {
        return OpenAiChatModel.builder()
                .apiKey(System.getenv("aliyunQwen-apiKey"))
                .modelName("qwen3-next-80b-a3b-instruct")
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .build();
    }

    /**
     * 基于低阶模型实例构建高阶使用的AIServices实例
     *
     * @param chatModel 对话模型
     * @return AIServices实例
     */
    @Bean(value = "chat")
    public ChatAssistant chatAssistant(ChatModel chatModel) {
        return AiServices.create(ChatAssistant.class, chatModel);
    }

    /**
     * 构建高阶AIServices实例(基于保留消息数的滑动窗口)
     *
     * @param chatModel 对话模型
     * @return AIServices实例
     */
    @Bean(value = "chatMemoryWithMessageWindow")
    public ChatMemoryAssistant chatMemoryWithMessageWindow(ChatModel chatModel) {
        return AiServices.builder(ChatMemoryAssistant.class)
                .chatModel(chatModel)
                //根据memoryId构建一个ChatMemory，保留最多100条消息
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(100))
                .build();
    }

    /**
     * 构建高阶AIServices实例(基于保留Token数的滑动窗口)
     *
     * @param chatModel 对话模型
     * @return AIServices实例
     */
    @Bean(value = "chatMemoryWithTokenWindow")
    public ChatMemoryAssistant chatMemoryWithTokenWindow(ChatModel chatModel) {
        // 构建默认的token分词器
        TokenCountEstimator openAiTokenizer = new OpenAiTokenCountEstimator("gpt-4");
        return AiServices.builder(ChatMemoryAssistant.class)
                .chatModel(chatModel)
                //根据memoryId构建一个ChatMemory，保留最多100条消息
                .chatMemoryProvider(memoryId -> TokenWindowChatMemory.withMaxTokens(1000, openAiTokenizer))
                .build();
    }
}
